
import asyncio
import json
import logging
import os

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from paho.mqtt import client as mqtt_client

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Mount static files
frontend_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "frontend")
if os.path.exists(frontend_path):
    app.mount("/static", StaticFiles(directory=frontend_path), name="static")

# MQTT Broker settings
MQTT_BROKER = "111.229.153.118"  # Replace with your MQTT Broker address
MQTT_PORT = 1883
MQTT_TOPIC = "robot/hexapod/status"  # The topic your robot publishes to
MQTT_CLIENT_ID = "python_mqtt_backend"

# WebSocket connections
active_websocket_connections: list[WebSocket] = []

# MQTT Client setup
def connect_mqtt() -> mqtt_client.Client:
    def on_connect(client, userdata, flags, rc):
        if rc == 0:
            logger.info("Connected to MQTT Broker!")
        else:
            logger.error(f"Failed to connect, return code {rc}\n")

    client = mqtt_client.Client(mqtt_client.CallbackAPIVersion.VERSION1, MQTT_CLIENT_ID)
    client.username_pw_set("admin", "public")
    client.on_connect = on_connect
    client.connect(MQTT_BROKER, MQTT_PORT)
    return client

mqtt_client_instance = connect_mqtt()

def subscribe_mqtt_topic():
    def on_message(client, userdata, msg):
        logger.info(f"Received `{msg.payload.decode()}` from `{msg.topic}` topic")
        message = msg.payload.decode()
        # Broadcast message to all active WebSocket connections
        asyncio.run(broadcast_message(message))

    mqtt_client_instance.subscribe(MQTT_TOPIC)
    mqtt_client_instance.on_message = on_message

async def broadcast_message(message: str):
    for connection in active_websocket_connections:
        try:
            await connection.send_text(message)
        except RuntimeError as e:
            logger.error(f"Error sending message to WebSocket: {e}")
        except WebSocketDisconnect:
            logger.warning("WebSocket disconnected during broadcast.")
            active_websocket_connections.remove(connection)

@app.on_event("startup")
async def startup_event():
    # Start MQTT loop in a non-blocking way
    mqtt_client_instance.loop_start()
    subscribe_mqtt_topic()
    logger.info("Backend service started.")

@app.on_event("shutdown")
async def shutdown_event():
    mqtt_client_instance.loop_stop()
    mqtt_client_instance.disconnect()
    logger.info("Backend service shut down.")

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    active_websocket_connections.append(websocket)
    logger.info(f"New WebSocket connection: {websocket.client}")
    try:
        while True:
            # Keep the connection alive, or handle incoming messages from frontend if any
            await websocket.receive_text()  # This will block until a message is received or connection closes
    except WebSocketDisconnect:
        active_websocket_connections.remove(websocket)
        logger.info(f"WebSocket disconnected: {websocket.client}")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")

@app.get("/")
async def get():
    # Serve the frontend HTML file
    frontend_file = os.path.join(frontend_path, "index.html")
    if os.path.exists(frontend_file):
        with open(frontend_file, "r", encoding="utf-8") as f:
            return HTMLResponse(content=f.read())
    else:
        return HTMLResponse("<h1>六足机器人监控后端服务</h1><p>WebSocket endpoint: /ws</p><p>前端文件未找到</p>")


